import cv2
import numpy as np
import os
import math

'''水平投影'''
thre_char = 100
def file_name(file_dir):
    alla = []
    for root, dirs, files in os.walk(file_dir):
        alla=files
    return alla
def local_file_name(file_dir):
    alla = []
    for root, dirs, files in os.walk(file_dir):
        alla=files
    return alla

def mkdir(path):
    folder = os.path.exists(path)

    if not folder:  # 判断是否存在文件夹如果不存在则创建为文件夹
        os.makedirs(path)  # makedirs 创建文件时如果路径不存在会创建这个路径
    else:
        pass

def getHProjection(image):
    hProjection = np.zeros(image.shape, np.uint8)
    # 图像高与宽
    (h, w) = image.shape
    # 长度与图像高度一致的数组
    h_ = [0] * h
    # 循环统计每一行白色像素的个数
    for y in range(h):
        for x in range(w):
            if image[y, x] == 255:
                h_[y] += 1
    # 绘制水平投影图像
    for y in range(h):
        for x in range(h_[y]):
            hProjection[y, x] = 255
    #cv2.imshow('hProjection2', hProjection)

    return h_


def getVProjection(image):
    vProjection = np.zeros(image.shape, np.uint8);
    # 图像高与宽
    (h, w) = image.shape
    # 长度与图像宽度一致的数组
    w_ = [0] * w
    # 循环统计每一列白色像素的个数
    for x in range(w):
        for y in range(h):
            if image[y, x] == 255:
                w_[x] += 1
    # 绘制垂直平投影图像
    for x in range(w):
        for y in range(h - w_[x], h):
            vProjection[y, x] = 255
    # cv2.imshow('vProjection',vProjection)
    return w_


def form(order,file_name):
#传入文件名，例如“Task1\\Task1_Test\\test-001.tiff”

    # 读入原始图像
#C:\\Users\\Robert\\Desktop\\SRTP
    origineImage = cv2.imread(file_name)
    # 图像灰度化
    #image = cv2.imread('test.jpg',0)
    image = cv2.cvtColor(origineImage, cv2.COLOR_BGR2GRAY)
    #cv2.imshow('gray', image)
    # 将图片二值化
    retval, img = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY_INV)
    #cv2.imshow('binary', img)
    # 图像高与宽
    (h, w) = img.shape
    Position = []
    # 水平投影
    H = getHProjection(img)
 #   print(len(H),(h,w))

    start = 0
    H_Start = []
    H_End = []
    # 根据水平投影获取垂直分割位置
    flag = 0
    for i in range(len(H)):
        if H[i] > 1 and start == 0:
            flag = i
            H_Start.append(i)
            start = 1
        if H[i] <= 1 and start == 1 and (i-flag)>thre_char:
            H_End.append(i)
            start = 0
    # 分割行，分割之后再进行列分割并保存分割位置
    #print(H,H_Start,H_End)
    for i in range(min(len(H_End),len(H_Start))):
        # 获取行图像
        cropImg = img[H_Start[i]:H_End[i], 0:w]
        # cv2.imshow('cropImg',cropImg)
        # 对行图像进行垂直投影
        W = getVProjection(cropImg)
       # print(W)
        Wstart = 0
        Wend = 0
        W_Start = 0
        W_End = 0
        for j in range(len(W)):
            if W[j] > 0 and Wstart == 0:
                W_Start = j
                Wstart = 1
                Wend = 0
            if W[j] <= 0 and Wstart == 1:
                W_End = j
                Wstart = 0
                Wend = 1
            if Wend == 1:
                Position.append([W_Start, H_Start[i], W_End, H_End[i]])
                Wend = 0
    # 根据确定的位置分割字符



    mkdir("dataset/test/person-"+str(order))
    mkdir("dataset/train/person-"+str(order))
    mkdir("dataset/valid/person-"+str(order))

    Name_test = local_file_name("dataset/test/person-"+str(order))
    test_amt=len(Name_test)+1
    Name_train = local_file_name("dataset/train/person-"+str(order))
    train_amt = len(Name_train)+1
    Name_valid = local_file_name("dataset/valid/person-"+str(order))
    valid_amt = len(Name_valid)+1

    sum_area=0
    thre_sum = 0

    for m in range(len(Position)):
        sum_area=sum_area+(Position[m][3]-Position[m][1])*(Position[m][2]-Position[m][0])
        thre_sum+=img[Position[m][1]:Position[m][3], Position[m][0]:Position[m][2]].sum()
    mean_area=sum_area/len(Position)
    thre_sum = thre_sum/len(Position)
    for m in range(len(Position)):

        if(Position[m][3]-Position[m][1])*(Position[m][2]-Position[m][0])<mean_area/2 or img[Position[m][1]:Position[m][3], Position[m][0]:Position[m][2]].sum()<thre_sum:
            continue

        img1 = cv2.resize(img[Position[m][1]:Position[m][3], Position[m][0]:Position[m][2]],(224, 224))
        if m<len(Position)/6:

            cv2.imwrite("dataset/test/person-"+str(order)+"/{}-No.".format(order) + "{}.png".format(test_amt), img1)
            test_amt=test_amt+1
        else:
            cv2.imwrite("dataset/train/person-"+str(order)+"/{}-No.".format(order) +"{}.png".format(train_amt),img1)
            train_amt=train_amt+1

